#include "gbn_cuda.cuh"
#include <iostream>
#include <random>
#include <chrono>
#include <opencv2/opencv.hpp>  // 添加OpenCV头文件
using namespace cv;

// 打印程序帮助信息（支持2D/高维场景配置说明）
void printUsage() {
    std::cout << "GBN Point Set Optimizer (2D/High-Dimensional Support)\n";
    std::cout << "Usage: gbn_optimizer [options]\n";
    std::cout << "Options:\n";
    std::cout << "  -n <num>   Number of points (default: 4096)\n";
    std::cout << "  -d <dim>   Dimensions (1~32, default: 2)\n";
    std::cout << "  -w <width> 2D Domain width (default: 100.0)\n";
    std::cout << "  -h <height> 2D Domain height (default: 100.0)\n";
    std::cout << "  -s <sigma> Gaussian kernel sigma (default: 1.0)\n";
    std::cout << "  -lr <rate> Learning rate (default: 0.1)\n";
    std::cout << "  -i <iter>  Max iterations (default: 10000)\n";
    std::cout << "  -t         Disable toroidal domain (default: enabled)\n";
    std::cout << "  -snap <path> Snapshot output path (default: ./gbn_snapshots/)\n";
    std::cout << "  -bins <num> Power spectrum bins (default: 100)\n";
    std::cout << "  -help      Show this help message\n";
}

// 主函数（模板化适配Float类型，默认float）
template <typename Float = float>
int runOptimizer(int argc, char** argv) {
    try {
        // -------------------------- 1. 初始化默认参数（适配2D/高维） --------------------------
        int num_points = 4096;                // 点集数量4K
        int dimensions = 2;                   // 默认2D（支持1~32维）
        std::vector<Float> domain_ranges;     // 各维度域范围（2D默认{1.0, 1.0}）
        Float sigma = DEFAULT_SIGMA<Float>;   // 高斯核带宽（派生类默认1.0）
        Float learning_rate = 0.1f;           // 学习率（派生类默认0.1）
        int max_iterations = 10000;           // 最大迭代次数（派生类默认10000）
        bool is_toroidal = true;              // 环形域（默认启用）
        std::string snap_path = SNAPSHOT_BASE_PATH; // 快照输出路径
        int power_spectrum_bins = 100;        // 功率谱区间数

        // 2D默认域范围初始化
        domain_ranges.push_back(1.0f); // 宽度
        domain_ranges.push_back(1.0f); // 高度

        // -------------------------- 2. 解析命令行参数（覆盖默认值） --------------------------
        for (int i = 1; i < argc; ++i) {
            std::string arg = argv[i];
            if (arg == "-help") {
                printUsage();
                return 0;
            }
            // 点集数量
            else if (arg == "-n" && i + 1 < argc) {
                num_points = std::stoi(argv[++i]);
                if (num_points <= 0) throw std::invalid_argument("Number of points must be positive");
            }
            // 维度（1~32，匹配基类MAX_DIMENSIONS）
            else if (arg == "-d" && i + 1 < argc) {
                dimensions = std::stoi(argv[++i]);
                if (dimensions < 1 || dimensions > MAX_DIMENSIONS) {
                    throw std::invalid_argument(
                        "Dimensions must be 1~" + std::to_string(MAX_DIMENSIONS) + 
                        " (current: " + std::to_string(dimensions) + ")"
                    );
                }
                // 高维场景：更新域范围（各维度默认100.0）
                if (dimensions != 2) {
                    domain_ranges.resize(dimensions, 100.0f);
                    std::cout << "[Info] High-dimensional mode (" << dimensions << "D): Domain range per dim = 100.0\n";
                }
            }
            // 2D宽度（仅2D生效）
            else if (arg == "-w" && i + 1 < argc && dimensions == 2) {
                domain_ranges[0] = std::stof(argv[++i]);
                if (domain_ranges[0] <= 0) throw std::invalid_argument("2D width must be positive");
            }
            // 2D高度（仅2D生效）
            else if (arg == "-h" && i + 1 < argc && dimensions == 2) {
                domain_ranges[1] = std::stof(argv[++i]);
                if (domain_ranges[1] <= 0) throw std::invalid_argument("2D height must be positive");
            }
            // 高斯核带宽
            else if (arg == "-s" && i + 1 < argc) {
                sigma = std::stof(argv[++i]);
                if (sigma <= 0) throw std::invalid_argument("Sigma must be positive");
            }
            // 学习率
            else if (arg == "-lr" && i + 1 < argc) {
                learning_rate = std::stof(argv[++i]);
                if (learning_rate <= 0) throw std::invalid_argument("Learning rate must be positive");
            }
            // 最大迭代次数
            else if (arg == "-i" && i + 1 < argc) {
                max_iterations = std::stoi(argv[++i]);
                if (max_iterations <= 0) throw std::invalid_argument("Max iterations must be positive");
            }
            // 禁用环形域
            else if (arg == "-t") {
                is_toroidal = false;
            }
            // 快照路径
            else if (arg == "-snap" && i + 1 < argc) {
                snap_path = argv[++i];
            }
            // 功率谱区间数
            else if (arg == "-bins" && i + 1 < argc) {
                power_spectrum_bins = std::stoi(argv[++i]);
                if (power_spectrum_bins <= 0) throw std::invalid_argument("Power spectrum bins must be positive");
            }
            // 未知参数
            else {
                throw std::invalid_argument("Unknown option: " + arg + " (use -help for usage)");
            }
        }

        // -------------------------- 3. 创建GBNPointSet实例（派生类核心逻辑） --------------------------
        // 初始化派生类：传入点数量+域范围（基类构造自动处理GPU缓冲区）
        GBNPointSet<Float> point_set(num_points, domain_ranges);
        
        // 配置优化参数（派生类成员）
        point_set.sigma = sigma;
        point_set.learning_rate = learning_rate;
        point_set.max_iterations = max_iterations;
        
        // 配置域参数（基类接口）
        point_set.setToroidalMode(is_toroidal);
        
        // 配置快照（基类默认参数：2D多格式，高维仅TXT；按2的幂次触发）
        point_set.setFrameBaseName(snap_path);
        point_set.initDefaultSnapshotParams(); // 基类自动适配2D/高维快照格式

        // -------------------------- 4. 执行优化（派生类核心接口） --------------------------
        std::cout << "================================ GBN Optimization Start ================================\n";
        std::cout << "Config: N=" << num_points << ", D=" << dimensions << ", Sigma=" << sigma 
                  << ", LR=" << learning_rate << ", Iter=" << max_iterations 
                  << ", Toroidal=" << (is_toroidal ? "Enabled" : "Disabled") << "\n";
        std::cout << "Snapshot Path: " << snap_path << "\n";
        
        auto start_time = std::chrono::high_resolution_clock::now();
        point_set.optimize(); // 派生类内化GPU梯度下降逻辑
        auto end_time = std::chrono::high_resolution_clock::now();
        
        // 计算优化耗时
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);
        std::cout << "================================ GBN Optimization End ================================\n";
        std::cout << "Optimization Time: " << duration.count() << " ms\n";

        // -------------------------- 5. 结果验证与可视化（基类接口） --------------------------
        
        // 5.1 计算平均最近邻距离（基类辅助接口：评估分布均匀性）
        Float avg_dist = point_set.calcAverageDist();
        std::cout << "Normalized Average Nearest Neighbor Distance: " << avg_dist << "\n";

        // 5.2 计算径向功率谱（基类接口：2D/高维均支持）
        std::vector<Float> power_spectrum;
        point_set.calcRadialPowerSpectrum(power_spectrum, power_spectrum_bins);
        
        // 5.3 保存功率谱（基类接口：2D支持PNG/EPS/PDF，高维自动跳过）
        std::string spectrum_path = snap_path + "power_spectrum.png";
        if (dimensions == 2) {
            point_set.saveRadialPowerSpectrum(power_spectrum, spectrum_path, 800); // 2D保存PNG
            point_set.showRadialPowerSpectrum(power_spectrum, "Optimized Power Spectrum", 800); // 2D显示
        }
        else {
            std::cout << "[Info] High-dimensional mode: Power spectrum visualization disabled (only TXT saved)\n";
        }

        // 5.4 输出低频功率（蓝噪声特征：低频功率低且均匀）
        std::cout << "Low-Frequency Power (First 10 Bins):\n";
        int print_bins = std::min(10, power_spectrum_bins);
        for (int i = 0; i < print_bins; ++i) {
            std::cout << "Bin " << i << ": " << power_spectrum[i] << "\n";
        }

        // -------------------------- 6. 保存最终点集（基类接口） --------------------------
        std::string final_points_path = snap_path + "final_points_" + std::to_string(dimensions) + "d.txt";
        point_set.saveTXT(final_points_path); // 基类TXT保存（2D/高维均支持）
        std::cout << "Final Points Saved To: " << final_points_path << "\n";

        // -------------------------- 7. 2D场景额外可视化（基类接口） --------------------------
        if (dimensions == 2) {
            // 保存最终点集PNG（2D矢量图）
            std::string final_png_path = snap_path + "final_points_2d.png";
            point_set.savePNG(final_png_path);
            std::cout << "Final 2D Points PNG Saved To: " << final_png_path << "\n";

            // 等待用户关闭窗口（基类show接口已处理按键，此处仅2D需等待）
            std::cout << "Press any key to close windows...\n";
            cv::waitKey(0);
            cv::destroyAllWindows();
        }

    } catch (const std::exception& e) {
        std::cerr << "\nError: " << e.what() << "\n";
        printUsage();
        return 1;
    }

    return 0;
}

// 主函数入口（默认float类型，兼顾精度与GPU性能）
int main(int argc, char** argv) {
    return runOptimizer<float>(argc, argv);
}